# -*- coding: utf-8 -*-

import cv2
import numpy as np
import os
import copy
import matplotlib.pyplot as plt
from PIL import Image

# 一开始是扩增方法
# 调用方法在最下面
# 这个好像是用于图像分割，所以有个mask

# 椒盐噪声
def SaltAndPepper(src, percetage):
    SP_NoiseImg = src.copy()
    SP_NoiseNum = int(percetage * src.shape[0] * src.shape[1])
    for i in range(SP_NoiseNum):
        randR = np.random.randint(0, src.shape[0] - 1)
        randG = np.random.randint(0, src.shape[1] - 1)
        randB = np.random.randint(0, 3)
        if np.random.randint(0, 1) == 0:
            SP_NoiseImg[randR, randG, randB] = 0
        else:
            SP_NoiseImg[randR, randG, randB] = 255
    return SP_NoiseImg


# 高斯噪声
# percetage代表的是出现的噪声点数量，越高，噪声点越多
def addGaussianNoise(image, percetage):
    G_Noiseimg = image.copy()
    w = image.shape[1]
    h = image.shape[0]
    G_NoiseNum = int(percetage * image.shape[0] * image.shape[1])
    for i in range(G_NoiseNum):
        temp_x = np.random.randint(0, h)
        temp_y = np.random.randint(0, w)
        G_Noiseimg[temp_x][temp_y][np.random.randint(3)] = np.random.randn(1)[0]
    return G_Noiseimg


# 昏暗
def darker(image, percetage=0.9):
    image_copy = image.copy()
    w = image.shape[1]
    h = image.shape[0]
    # get darker
    for xi in range(0, w):
        for xj in range(0, h):
            image_copy[xj, xi, 0] = int(image[xj, xi, 0] * percetage)
            image_copy[xj, xi, 1] = int(image[xj, xi, 1] * percetage)
            image_copy[xj, xi, 2] = int(image[xj, xi, 2] * percetage)
    return image_copy


# 亮度
def brighter(image, percetage=1.5):
    image_copy = image.copy()
    w = image.shape[1]
    h = image.shape[0]
    # get brighter
    for xi in range(0, w):
        for xj in range(0, h):
            image_copy[xj, xi, 0] = np.clip(int(image[xj, xi, 0] * percetage), a_max=255, a_min=0)
            image_copy[xj, xi, 1] = np.clip(int(image[xj, xi, 1] * percetage), a_max=255, a_min=0)
            image_copy[xj, xi, 2] = np.clip(int(image[xj, xi, 2] * percetage), a_max=255, a_min=0)
    return image_copy


# 旋转 angle度数
def rotate(image, angle, center=None, scale=1.0):
    (h, w) = image.shape[:2]
    # If no rotation center is specified, the center of the image is set as the rotation center
    if center is None:
        center = (w / 2, h / 2)
    m = cv2.getRotationMatrix2D(center, angle, scale)
    rotated = cv2.warpAffine(image, m, (w, h), borderValue=(255, 255, 255))
    return rotated


# 翻转
def flip(image):
    flipped_image = np.fliplr(image)
    return flipped_image


# 颜色反转
def image_inverse(input):
    value_max = np.max(input)
    output = value_max - input
    return output


# 这个没用，是将数据都缩小到一定范围内，较少了像素之间的区别，不适合分割任务
def image_log(input):
    output = np.copy(input)
    output[..., 0] = np.log(input[..., 0] + 1)
    output[..., 1] = np.log(input[..., 1] + 1)
    output[..., 2] = np.log(input[..., 2] + 1)

    return output


# 增加亮度 或者是提升对比度
def img_conv(input, a=1.5, b=0):
    output = cv2.convertScaleAbs(input, alpha=a, beta=b)
    return output


# 限制阈值的自适应直方图均衡化
# clipLimit未像素值阈值，tileGridSize在多大范围内进行自适应均衡化
def equ_patch(input, clipLimit=190, tileGridSize=(8, 8)):
    clahe = cv2.createCLAHE(clipLimit, tileGridSize)
    output = clahe.apply(input)
    return output


# 全局均衡化 但是两个均衡化需要将数据规范到uin8上，线性转化 后期有需要进行修改
def equa_global(input):
    out = cv2.equalizeHist(input)
    return out


def img_resize(input, h, w):
    out = cv2.resize(input, (h, w))
    return out


# 这里区分是否是mask的作用是，针对mask不做加噪处理，只是对img进行加噪
# 同时在进行实际的扩增时，自己展示一下图片，是否是自己想要的，确认好再进行扩增
def Argu1(img, mask=False):
    if mask == False:  # 使用旋转 椒盐噪声 提升亮度
        img_argu = rotate(img, 20)
        img_argu = SaltAndPepper(img_argu, 0.7)
        img_argu = brighter(img_argu, 1.3)
    else:
        img_argu = rotate(img, 20)
        img_argu = brighter(img_argu, 1.3)
        # img_argu = cv2.resize(img_argu, (256, 256))
    return img_argu


def Argu2(img, mask=False):
    if mask == False:  # 使用旋转 高斯噪声 降低亮度
        img_argu = brighter(img, 0.9)
        img_argu = rotate(img_argu, -20)
        img_argu = addGaussianNoise(img_argu, 0.7)
    else:
        img_argu = brighter(img, 1)
        img_argu = rotate(img_argu, -20)
    # img_argu = cv2.resize(img_argu, (256, 256))
    return img_argu

# 调用数据扩增方法
def ArguLihang(img,img_path,imagePath,resultPath,filename):
    saveImg(img, img_path, imagePath, resultPath)
    imgFlip=flip(img)
    newImg_path=img_path.replace(filename.split("_")[1],"flip_"+filename.split("_")[1])
    saveImg(imgFlip,newImg_path,imagePath,resultPath)
    for i in range(1,17):
        rote=i*20
        img_argu = rotate(img, rote)
        newImg_path = img_path.replace(filename.split("_")[1], "rotate" +str(rote)+"_"+ filename.split("_")[1])
        saveImg(img_argu, newImg_path, imagePath, resultPath)
    #img_argu = rotate(img, -20)


# 保存文件方法  TODO   获取文件名，对文件名根据扩增方式进行修改
def saveImg(img,img_path,imagePath,resultPath):
    newImgPath= img_path.replace(imagePath,resultPath)
    dir_name = os.path.dirname(newImgPath)
    checkDirForDeepIsExists(dir_name)
    cv2.imwrite(newImgPath, img)

def checkDirForDeepIsExists(dir_name):
    if not os.path.exists(dir_name):  # os模块判断并创建
        deep_dir_name = os.path.dirname(dir_name)
        checkDirForDeepIsExists(deep_dir_name)
        os.mkdir(dir_name)

from tqdm import tqdm

# 如果结果文件夹是原文件夹的情况可以在ArguLihang中把保存原图给去了
imagePath="H:\\test\\before"
resultPath="H:\\test\\after"
for dirpath, dirnames, filenames in tqdm(os.walk(imagePath)):
    for filename in filenames:
        img_path=os.path.join(dirpath, filename)
        #print(os.path.join(dirpath, filename))
        img = cv2.imread(img_path)
        ArguLihang(img,img_path,imagePath,resultPath,filename)

# 原版
# file_dir = r'E:/Seg/data/imgs/'  # 这里是需要扩增的文件夹目录11，后面加上/
# file_argu = r'E:/Seg/data/imgs/'  # 这里是建立在同一目录下的扩增文件夹目录22， 后面加上/
# for img_name in tqdm(os.listdir(file_dir)):  # tqdm是用来展示遍历进程的， 同时获取到11中的文件，遍历，读取，转换
#     img_path = file_dir + img_name
#     img = cv2.imread(img_path)  # 这里默认分割数据集img和mask名称，除了后缀不同，其余都相同
#     mask = cv2.imread(img_path.replace("imgs", "masks").replace("jpg", "tif").replace("bmp", "tif"))  # 读取对应的mask文件
#
#     img_1 = Argu1(img, False)
#     mask_1 = Argu1(mask, True)
#     img_2 = Argu2(img, False)
#     mask_2 = Argu2(mask, True)
#
#     cv2.imwrite(file_argu + img_name[0:-4] + '.jpg', img)
#     cv2.imwrite(file_argu.replace("imgs", "masks") + img_name[0:-4] + '.tif', mask)
#
#     cv2.imwrite(file_argu + img_name[0:-4] + '_Arg1.jpg', img_1)  # 扩增后的图片，在图片名称后加一个标识用来和原始图片进行区分
#     cv2.imwrite(file_argu.replace("imgs", "masks") + img_name[0:-4] + '_Arg1.jpg'.replace("jpg", "tif"), mask_1)
#
#     cv2.imwrite(file_argu + img_name[0:-4] + '_Arg2.jpg', img_2)
#     cv2.imwrite(file_argu.replace("imgs", "masks") + img_name[0:-4] + '_Arg2.jpg'.replace("jpg", "tif"), mask_2)

